import time
import numpy as np
import matplotlib.pyplot as plt
from path_planning.a_star import AStarPlanner
from path_planning.hybrid_a_star import HybridAStarPlanner
from path_planning.dwa_planner import DWAPlanner
from utils.grid_map import GridMap
from visualization.plotter import PathPlotter

def main():
    # Configuration
    grid_size = (100, 100)  # 100x100 grid
    resolution = 1.0  # 1m per grid cell
    obstacle_density = 0.15  # 15% obstacle density
    start = [5.0, 5.0, np.pi/4]  # [x, y, theta] in world coordinates
    goal = [90.0, 90.0]  # [x, y] in world coordinates

    # Create grid map with obstacles
    print("Generating grid map...")
    grid_map = GridMap(size=grid_size, resolution=resolution)
    grid_map.generate_random_obstacles(obstacle_density)

    # Add dynamic obstacles
    dynamic_obstacle_ids = []
    dynamic_obstacle_ids.append(grid_map.add_dynamic_obstacle([20, 20], [80, 20], speed=0.8))
    dynamic_obstacle_ids.append(grid_map.add_dynamic_obstacle([20, 40], [80, 40], speed=0.6))
    dynamic_obstacle_ids.append(grid_map.add_dynamic_obstacle([80, 60], [20, 60], speed=0.7))

    # Initialize plotter
    plotter = PathPlotter(grid_map, resolution)

    # Run A* algorithm
    print("Running A* algorithm...")
    a_star_planner = AStarPlanner(grid_map.get_grid(), resolution)
    start_time = time.time()
    a_star_path = a_star_planner.planning(start[:2], goal)
    a_star_time = (time.time() - start_time) * 1000  # Convert to milliseconds
    print(f"A* planning completed in {a_star_time:.2f} ms")

    # Run Hybrid A* algorithm
    print("Running Hybrid A* algorithm...")
    hybrid_a_star_planner = HybridAStarPlanner(grid_map.get_grid(), resolution, min_turning_radius=6.0)
    start_time = time.time()
    hybrid_a_star_path = hybrid_a_star_planner.planning(start, goal)
    hybrid_a_star_time = (time.time() - start_time) * 1000
    print(f"Hybrid A* planning completed in {hybrid_a_star_time:.2f} ms")

    # Initialize hybrid path simplified
    hybrid_path_simplified = []
    
    # Plot results
    if a_star_path:
        plotter.plot_path(a_star_path, color='r', label=f'A* Path ({a_star_time:.2f} ms)')
    if hybrid_a_star_path:
        # Convert Hybrid A* path format to match A* path format for plotting
        hybrid_path_simplified = [[p[0], p[1]] for p in hybrid_a_star_path]
        plotter.plot_path(hybrid_path_simplified, color='g', label=f'Hybrid A* Path ({hybrid_a_star_time:.2f} ms)')
    else:
        print("Hybrid A* did not find a path")

    # Plot start and goal
    plotter.ax.plot(start[0], start[1], 'go', markersize=12, label='Start')
    plotter.ax.plot(goal[0], goal[1], 'bo', markersize=12, label='Goal')

    # Plot dynamic obstacles
    dynamic_obs_positions = [(grid_map.grid2world([obs['current'][0], obs['current'][1]])) for obs in grid_map.dynamic_obstacles]
    plotter.plot_dynamic_obstacles(dynamic_obs_positions)

    # Show plot
    plotter.show(block=False)

    # DWA local planner simulation
    print("Running DWA local planner simulation...")
    dwa_planner = DWAPlanner(max_speed=2.0, max_accel=0.5, max_yawrate=40.0*np.pi/180.0)
    current_state = [start[0], start[1], start[2], 0.0, 0.0]  # [x, y, yaw, v, yawrate]
    dwa_trajectory = [current_state[:3]]
    simulation_time = 0
    max_simulation_time = 100  # seconds

    # Create DWA visualization
    dwa_plotter = PathPlotter(grid_map, resolution)
    dwa_plotter.plot_path(hybrid_path_simplified, color='g', linewidth=1, label='Global Path')

    while simulation_time < max_simulation_time:
        # Update dynamic obstacles
        dyn_obs_pos = grid_map.update_dynamic_obstacles(dt=0.1)
        dyn_obs_world = [grid_map.grid2world(pos) for pos in dyn_obs_pos]

        # Run DWA planning
        best_u, trajectory = dwa_planner.plan(current_state, goal, grid_map, dyn_obs_world)

        if not best_u or not trajectory:
            print("No valid trajectory found. Stopping simulation.")
            break

        # Update current state
        current_state[3] = best_u[0]  # velocity
        current_state[4] = best_u[1]  # yaw rate

        # Update position using bicycle model
        dt = 0.1
        current_state[0] += current_state[3] * np.cos(current_state[2]) * dt
        current_state[1] += current_state[3] * np.sin(current_state[2]) * dt
        current_state[2] += current_state[4] * dt

        # Add to trajectory
        dwa_trajectory.append(current_state[:3])

        # Update visualization
        dwa_plotter.clear()
        dwa_plotter.plot_path(hybrid_path_simplified, color='g', linewidth=1, label='Global Path')
        dwa_plotter.plot_path(dwa_trajectory, color='b', linewidth=2, label='DWA Trajectory')
        dwa_plotter.plot_vehicle(current_state[0], current_state[1], current_state[2])
        dwa_plotter.plot_dynamic_obstacles(dyn_obs_world)
        dwa_plotter.ax.plot(goal[0], goal[1], 'bo', markersize=12, label='Goal')
        dwa_plotter.ax.set_title(f'DWA Local Planner Simulation (Time: {simulation_time:.1f}s)')
        dwa_plotter.ax.legend()
        plt.pause(0.01)

        # Check if goal is reached
        dist_to_goal = np.hypot(current_state[0] - goal[0], current_state[1] - goal[1])
        if dist_to_goal < 1.0:
            print("Goal reached!")
            break

        simulation_time += dt

    print("Simulation completed.")
    plt.show()

if __name__ == "__main__":
    main()